import os
import random
import shutil
path='./data/'
cate_list = os.listdir(path)  # 获取未分词语料库中所有类别
for mydir in cate_list:
    train_dir = './train_corpus/' + mydir + "/"  # 拼出训练集存贮的对应目录路径如：train_corpus/体育/
    if not os.path.exists(train_dir):  # 是否存在分词目录，如果没有则创建该目录
        os.makedirs(train_dir)

    test_dir = './test_corpus/' + mydir+'/' # 拼出测试存贮的对应目录路径如：test_corpus/体育/
    if not os.path.exists(test_dir):
        os.makedirs(test_dir)
    class_path=path+mydir+'/' # 生成当前类别的目录 如 data/体育/
    file_list=os.listdir(class_path) # 生成分类下文本名称列表
    length=len(file_list)
    print(mydir,length)
    res = random.sample(range(1, length+1), 5000) #在所有文件中随机选取五千个生成索引列表,乱序
    train=res[0:4000]#训练集占80%，4000条
    test=res[4000:5000]#测试集占20%，1000条
    #print(len(train),len(test))
    for file_path in file_list:
        fullname = class_path + file_path  # 拼出文件名全路径如：data/体育/21.txt
        #print(file_path)
        x=int(file_path.split('.')[0].split('_')[1])#将文件名中的序号取出
        if x in train:
           shutil.copyfile(fullname, train_dir+file_path)#将前面的文件复制到后面的目录中
        if x in test:
           shutil.copyfile(fullname, test_dir + file_path)

